#!/usr/bin/python

import anydbm
import cPickle
import twitter

class TweetStore(object):
  """A persistent store for all tweets (twitter.Status) messages"""

  def __init__(self,filename=None):
    if filename:
      self._in_memory=False
      self._store = anydbm.open(filename,'c')
    else:
      self._in_memory=True
      self._store = {}

  def __del__(self):
    self.close()

  def __itoa(self,int_id):
    return "%d" % int_id

  def __atoi(self,str_id): return int(str_id) 
  def __setitem__(self,key,value):
    self._store[self.__itoa(key)]=cPickle.dumps(value)

  def __getitem__(self,key):
    return cPickle.loads(self._store[self.__itoa(key)])

  def __delitem__(self,key):
    del self._store[self.__itoa(key)]

  def __contains__(self,key):
    if type(key) is twitter.Status:
      return self.__itoa(key.id) in self
    if type(key) is int:
      return self.__itoa(key) in self
    else:
      return key in self._store

  def keys(self):
    return [self.__atoi(keystr) for keystr in self._store.keys()]

  def values(self):
    return [cPickle.loads(valstr) for valstr in self._store.values()]

  def items(self):
    return [(self.__atoi(keystr),cPickle.loads(valstr)) for (keystr,valstr) in self._store.items()]

  def close(self):
    if not self._in_memory:
      self._store.close()

  def add(self,status):
    self[status.id] = status

  def matching(self,substr, ignore_case=False):
    """return tweets having text containing that substring."""
    if ignore_case: 
      filterfunc = lambda x: substr.lower() in x.text.lower()
    else:
      filterfunc = lambda x: substr in x.text

    return filter(filterfunc, self.values())

  def sync(self):
    if self._in_memory:
      self._store.sync()
    else:
      # nothing to do for an in-memory store
      pass

if __name__ == "__main__":

  import sys
  import codecs

  if len(sys.argv) > 1:
    dbfile = sys.argv[1]
    ts = TweetStore(dbfile)
    if len(sys.argv) > 2:
      # outfile = codecs.open(sys.argv[2],'w','utf-8')
      outfile = open(sys.argv[2],'wb')
    else:
      # this proved entirely ineffective.  
      # outfile = codecs.EncodedFile(sys.stdout, 'utf-8', file_encoding='ascii',errors='ignore')
      outfile = sys.stdout
      
    for tw in ts.values():
      text = filter(lambda x: x != '\n', tw.text)
      outfile.write(('%10d %20s: %s\n' % (tw.id, tw.user.screen_name, text)).encode('utf-8'))
    outfile.close()
    ts.close()
  else:
    sys.stderr.write('usage: %s <tweetdb_file> [<output_file>]\n' % sys.argv[0])
    sys.exit(1)

